import random
import sys
import time
from collections import defaultdict
from scapy.all import rdpcap
class CountMinSketch:
    def __init__(self, width, depth):
        self.width = width
        self.depth = depth
        self.sketch = [[0] * width for _ in range(depth)]
    
    def update(self, flow_hash):
        for i in range(self.depth):
            hash_val = (flow_hash + i) % self.width
            self.sketch[i][hash_val] += 1
    
    def estimate(self, flow_hash):
        min_count = float('inf')
        for i in range(self.depth):
            hash_val = (flow_hash + i) % self.width
            min_count = min(min_count, self.sketch[i][hash_val])
        return min_count
def calculate_flow_frequency(dataset, width, depth):
    flow_counter = CountMinSketch(width, depth)
    for packet in dataset:
        if packet.haslayer("IP"):
            src_ip = packet.payload.src  # 获取源IP
            dst_ip = packet.payload.dst  # 获取目的IP
            flow_hash = hash((src_ip, dst_ip))  # 计算流的哈希值
            flow_counter.update(flow_hash)  # 更新Count-Min Sketch
    return flow_counter
pcap_file = "test2.pcap"
dataset = rdpcap(pcap_file)  # 使用rdpcap函数读取pcap文件
width = 500  # Count-Min Sketch的宽度
depth = 5  # Count-Min Sketch的深度
start_time = time.time()  # 记录开始时间
flow_frequency = calculate_flow_frequency(dataset, width, depth)
end_time = time.time()  # 记录结束时间
total_time = end_time - start_time  # 计算总时间
element_size = sys.getsizeof(0)  # 获取一个元素的大小
array_size = width * depth  # 计算数组的总元素个数
mem_size = element_size * array_size
with open("CM-K-result.txt", "w") as file:
  for packet in dataset:
    if packet.haslayer("IP")and packet.haslayer("UDP"):
        src_ip = packet.payload.src  # 获取源IP
        dst_ip = packet.payload.dst  # 获取目的IP
        flow_hash = hash((src_ip, dst_ip))  # 计算流的哈希值
        flow_estimate = flow_frequency.estimate(flow_hash)  # 估计流的频率
        file.write(f"源IP: {src_ip}, 频率估计值: {flow_estimate}\n")
        print(f"源IP: {src_ip}, 频率估计值: {flow_estimate}")
print(f"count-min sketch二维数组占用的内存大小: {mem_size} bytes")
print(f"计算的总时间: {total_time} seconds")
